from langchain.prompts import PromptTemplate
from langchain.prompts import FewShotPromptTemplate
from langchain.prompts.example_selector import SemanticSimilarityExampleSelector
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import DashScopeEmbeddings
from hospitals.models import DepartmentModel
from doctors.models import DoctorModel

def get_answer(user_input):
    # words = request.query_params.get('word', None)
    # if words is None:
    #     return Response({"message": "关键词不能为空"})
    # 项目中,存储mysql中
    department = DepartmentModel.objects.all()
    examples = [{"id": str(i.id), "name": i.name} for i in department]
    print(examples)
    print("-------------------")
    # 有id,对应特征描述,此处测试,应替换数据库中的数据
    examples_prompt = PromptTemplate(
        input_variables=["id", "features"],
        template="根据你描述的病情，建议你挂{name}，推荐的医生为",
    )
    # 语义相似性示例选择器
    examples_selector = SemanticSimilarityExampleSelector.from_examples(
        examples,
        DashScopeEmbeddings(),
        FAISS,
        k=1
    )
    fewshot_prompt = FewShotPromptTemplate(
        example_selector=examples_selector,
        example_prompt=examples_prompt,
        prefix="",
        # # 这是提示的后缀部分，它通常包含了输入和输出的格式说明。在这个例子中，后缀是“输入:{adjective}\n 输出”，意味着模型需要接收一个形容词（adjective）作为输入，并产生一个输出
        suffix="",
        input_variables=['user_input']
    )
    res = fewshot_prompt.format(user_input=user_input)
    print(res)

    doctor = DoctorModel.objects.all()
    examples1 = [{"id": str(i.id), "name": i.name} for i in doctor]
    examples_prompt1 = PromptTemplate(
        input_variables=["id", "features"],
        template="{name}",
    )
    examples_selector = SemanticSimilarityExampleSelector.from_examples(
        examples1,
        DashScopeEmbeddings(),
        FAISS,
        k=4
    )
    fewshot_prompt1 = FewShotPromptTemplate(
        example_selector=examples_selector,
        example_prompt=examples_prompt1,
        prefix="",
        # # 这是提示的后缀部分，它通常包含了输入和输出的格式说明。在这个例子中，后缀是“输入:{adjective}\n 输出”，意味着模型需要接收一个形容词（adjective）作为输入，并产生一个输出
        suffix="",
        input_variables=['user_input']
    )
    res1 = fewshot_prompt1.format(user_input=res)
    items = res1.split('\n\n')
    print(items)
    # 生成新id列表
    doctor_list = ','.join(items)
    return res + doctor_list




